# -*- coding: utf-8 -*-

########################
# Created on 2022/05
# @author: ts
########################
import os
import queue
import re
import threading
import time
import copy
import cypcap
import struct
import socket
from scapy.all import rdpcap

# from .EthMessageIntf import EthMessageintf


REIP = re.compile('^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$')

# DoIP Header Structure : <protocol version><inverse protocol version><payload type><payloadlength><payload>
# Payload format : <local ecu address> <optional: target ecu addres> <optional message ><ASRBISO><ASRBOEM>
PROTOCOL_VERSION = DOIP_PV = '02'
INVERSE_PROTOCOL_VERSION = DOIP_IPV = 'FD'

# Payload type definitions#
DOIP_GENERIC_NEGATIVE_ACKNOWLEDGE = DOIP_NARP = '0000'
DOIP_VEHICLE_ID_REQUEST = '0001'
DOIP_VEHICLE_ID_REQUEST_W_EID = '0002'
DOIP_VEHICLE_ID_REQUEST_W_VIN = '0003'
DOIP_VEHICLE_ANNOUNCEMENT_ID_RESPONSE = '0004'
# DOIP_ROUTING_ACTIVATION_REQUEST : <0005><sourceaddress><activation type><00000000>
DOIP_ROUTING_ACTIVATION_REQUEST = DOIP_RAR = '0005'
# Activation Type
DEFAULT_ACTIVATION = '00'
WWH_OBD_ACTIVATION = '01'
# 0x02-0xDF ISOSAE Reserved
CENTRAL_SECURITY_ACTIVATION = 'E0'
# 0xE1-0xFF OEM Specific
ACTIVATION_SPACE_RESERVED_BY_ISO = ASRBISO = '00000000'
# the code above is mandatory but has no use at the moment. ISOSAE Reserved
ACTIVATION_SPACE_RESERVED_BY_OEM = ASRBOEM = 'ffffffff'

DOIP_ROUTING_ACTIVATION_RESPONSE = '0006'
DOIP_ALIVE_CHECK_REQUEST = '0007'
DOIP_ALIVE_CHECK_RESPONSE = '0008'
# 0x009-0x4000: Reserved by ISO13400
DOIP_ENTITY_STATUS_REQUEST = '4001'
DOIP_ENTITY_STATUS_RESPONSE = '4002'
DOIP_DIAGNOSTIC_POWER_MODE_INFO_REQUEST = '4003'
DOIP_DIAGNOSTIC_POWER_MODE_INFO_RESPONSE = '4004'
# 0x4005-0x8000 Reserved by ISO13400
DOIP_DIAGNOSTIC_MESSAGE = DOIP_UDS = '8001'
DOIP_DIAGNOSTIC_POSITIVE_ACKNOWLEDGE = '8002'
DOIP_DIAGNOSTIC_NEGATIVE_ACKNOWLEDGE = '8003'
# 0x8004-0xEFFF Reserved by ISO13400
# 0xF000-0xFFFF Reserved for manufacturer-specific use


payloadTypeDescription = {
    int(DOIP_GENERIC_NEGATIVE_ACKNOWLEDGE): "Generic negative response",
    int(DOIP_VEHICLE_ID_REQUEST): "Vehicle ID request",
    int(DOIP_VEHICLE_ID_REQUEST_W_EID): "Vehicle ID request with EID",
    int(DOIP_VEHICLE_ID_REQUEST_W_VIN): "Vehicle ID request with VIN",
    int(DOIP_VEHICLE_ANNOUNCEMENT_ID_RESPONSE): "Vehicle announcement ID response",
    int(DOIP_ROUTING_ACTIVATION_REQUEST): "Routing activation request",
    int(DOIP_ROUTING_ACTIVATION_RESPONSE): "Routing activation response",
    int(DOIP_ALIVE_CHECK_REQUEST): "Alive check request",
    int(DOIP_ALIVE_CHECK_RESPONSE): "Alive check response",
    int(DOIP_ENTITY_STATUS_REQUEST): "Entity status request",
    int(DOIP_ENTITY_STATUS_RESPONSE): "Entity status response",
    int(DOIP_DIAGNOSTIC_POWER_MODE_INFO_REQUEST): "Diagnostic power mode info request",
    int(DOIP_DIAGNOSTIC_POWER_MODE_INFO_RESPONSE): "Power mode info response",
    int(DOIP_DIAGNOSTIC_MESSAGE): "Diagnostic message",
    int(DOIP_DIAGNOSTIC_POSITIVE_ACKNOWLEDGE): "Diagnostic positive acknowledge",
    int(DOIP_DIAGNOSTIC_NEGATIVE_ACKNOWLEDGE): "Diagnostic negative acknowledge",
}


class EthCapture(object):
    def __init__(self, channel) -> None:
        self.__quit = False
        self.__dataHandlers = list()
        self.__capThread = None
        self._onEthData = None
        self.channel = channel
        self.__refCount = 1

    # def OnEthData(channel, srcip, srcport, destip, destport, data) -> bool:
    #     pass

    @property
    def OnEthData(self):
        """
        修改方法为属性
        """
        return self._onEthData

    @OnEthData.setter
    def OnEthData(self, func):
        """
        注册数据处理回调
        """
        self._onEthData = func

    def _CaptureThread(self, channel) -> None:
        while not self.__quit:
            cap = CyPcap()
            print(f"Used Channel: {channel}")
            if cap.FindEth(channel):
                cap.Start()
                for ts, src_ip, src_port, dst_ip, dst_port, payload in cap.Data():
                    if self.__quit:
                        break
                    if self._onEthData:
                        self._onEthData(self.channel, ts, src_ip, src_port, dst_ip, dst_port, payload)
                    # for handler in self.__dataHandlers:
                    #     handler(ts, src_ip, src_port, dst_ip, dst_port, payload)
            cap.Destroy()
            time.sleep(5)

    def Start(self) -> None:
        if not self.__capThread:
            self.__capThread = threading.Thread(target=self._CaptureThread, args=(self.channel,))
        if self.__refCount <= 1:
            self.__capThread.start()
        else:
            self.__refCount += 1

    def Stop(self) -> None:
        self.__refCount -= 1
        if self.__refCount <= 1:
            self.__quit = True

    def AppendHandler(self, handler) -> None:
        self.__dataHandlers.append(handler)


class IPParser:
    ETH_MAC_LENGTH = 14  # 报文MAC头部长度
    IP_HEADER_LENGTH = 20  # 报文ip头部长度
    TRANS_CTRL_LENGTH = 20  # 报文传输控制头部长度

    @classmethod
    def ParseEthHeader(cls, ethHeader) -> dict:
        """
        ETH报文格式解析
        :param ethHeader:
        :return:
        """
        dst_eth = ":".join(
            ["0{0:X}".format(imac) if (len(hex(imac)) % 2) != 0 else "{0:X}".format(imac) for imac in ethHeader[:6]])
        src_eth = ":".join(
            ['0{0:X}'.format(imac) if len(hex(imac)) % 2 != 0 else '{0:X}'.format(imac) for imac in ethHeader[6:12]])
        typ_eth = "IPv4" if ethHeader[12:14] == b'\x08\x00' else str(ethHeader[12:14])

        return {
            'dst_eth': dst_eth,
            'src_eth': src_eth,
            'typ_eth': typ_eth
        }

    @classmethod
    def ParseIpHeader(cls, ipHeader, isSrcPkg=False) -> dict:
        """
        IP报文格式解析
        :param ipHeader:
        :return:
        """
        if isSrcPkg:
            ipHeader = ipHeader[cls.ETH_MAC_LENGTH: cls.ETH_MAC_LENGTH + cls.IP_HEADER_LENGTH]
        line1 = struct.unpack('>BBH', ipHeader[:4])  # 先按照8位、8位、16位解析
        ip_version = ipHeader[0] >> 4  # 通过右移4位获取高四位
        # 报文头部长度的单位是32位 即四个字节
        iph_length = (line1[0] & 15) * 4  # 与1111与运算获取低四位
        packet_length = line1[2]
        line3 = struct.unpack('>BBH', ipHeader[8: 12])
        TTL = line3[0]
        protocol = line3[1]
        iph_checksum = line3[2]
        line4 = struct.unpack('>4s', ipHeader[12: 16])
        src_ip = socket.inet_ntoa(line4[0])
        line5 = struct.unpack('>4s', ipHeader[16: 20])
        dst_ip = socket.inet_ntoa(line5[0])

        # 返回结果
        return {
            'ip_version': ip_version,  # ip_version ip版本
            'iph_length': iph_length,  # iph_length ip头部长度
            'packet_length': packet_length,  # packet_length 报文长度
            'TTL': TTL,  # TTL 报文寿命
            'protocol': protocol,  # protocol 协议号 1 ICMP协议 6 TCP协议 17 UDP协议
            'iph_checksum': iph_checksum,  # iph_checksum ip头部的校验和
            'src_ip': src_ip,  # src_ip 源ip
            'dst_ip': dst_ip  # dst_ip 目的ip
        }

    @classmethod
    def ParseTransHeader(cls, transControl, isSrcPkg=False) -> dict:
        """
        传输控制协议报文解析
        :param transControl:
        :return:
        """
        if isSrcPkg:
            transControl = transControl[
                           cls.ETH_MAC_LENGTH + cls.IP_HEADER_LENGTH: cls.ETH_MAC_LENGTH + cls.IP_HEADER_LENGTH + cls.TRANS_CTRL_LENGTH]
        dst_port = struct.unpack('>H', transControl[:2])
        src_port = struct.unpack('>H', transControl[2:4])

        flags = bytearray([transControl[12] & 0x0f, transControl[13]])

        return {
            'dst_port': dst_port[0],
            'src_port': src_port[0],
            'flags': flags
        }

    @classmethod
    def parse(cls, packet) -> dict:
        ip_header = packet[cls.ETH_MAC_LENGTH: cls.ETH_MAC_LENGTH + cls.IP_HEADER_LENGTH]
        return cls.ParseIpHeader(ip_header)


class CyPcap():
    def __init__(self) -> None:
        self.capDev = None
        self.channel = None
        self.isFile = False
        self.pcap = None

    def FindEth(self, capDev) -> bool:
        if not REIP.match(capDev) and os.path.exists(capDev):
            self.isFile = True
            self.channel = capDev
            return True
        devs = cypcap.findalldevs()
        print(f"Find device: {devs}")
        self.capDev = None
        for dev in devs:
            for ip in dev.addresses:
                if ip.addr[1][0] == capDev:
                    self.capDev = dev
                    self.channel = capDev
                    return True
        return False

    def FilterCondition(self, src_ip, dst_port, src_port=0) -> dict:
        self.condition = {"src_ip": src_ip, "dst_port": dst_port, "src_port": src_port}

    def Filter(self, pkthdr, data, channel) -> None:
        # 过滤掉ARP等报文
        ethHeader = IPParser.ParseEthHeader(data)
        if ethHeader["typ_eth"] != "IPv4":
            # print(f"Unmatch ethernet type: IPv4")
            return
        ipHeader = IPParser.ParseIpHeader(data, True)
        if ipHeader["src_ip"] != self.condition["src_ip"]:
            # print(f"Unmatch src port: {ipHeader['src_ip']} and {self.condition['src_ip']}")
            return
        ports = IPParser.ParseTransHeader(data, True)
        if ports["dst_port"] != self.condition["dst_port"]:
            print(f"Unmatch dst port: {ports['dst_port']} and {self.condition['dst_port']}")
            return
        if self.condition["src_port"] and ports["src_port"] != self.condition["src_port"]:
            print(f"Unmatch src port: {ports['src_port']} and {self.condition['src_port']}")
            return
        doipPkg = data[IPParser.ETH_MAC_LENGTH + IPParser.IP_HEADER_LENGTH + IPParser.TRANS_CTRL_LENGTH:]
        # 忽略空包数据
        if not doipPkg:
            return
        print(f"doipPkg: {doipPkg}")
        # doipMsg = DoIPMsg(doipPkg)
        # print(doipMsg.PrintMessage())

    def Start(self) -> bool:
        if self.isFile:
            self.pcap = rdpcap("./normal.pcapng")
            return True
        try:
            self.pcap = cypcap.create(self.capDev)
            self.pcap.set_snaplen(65536)
            self.pcap.set_promisc(True)
            self.pcap.set_timeout(1000)
            self.pcap.activate()
            return True
        except Exception as e:
            print(f"Ethernet capture: {e}")
        return False

    def _waper_data(self):
        if self.isFile:
            for packet in self.pcap:
                packet.ts = packet.time
                yield packet, bytes(packet)
        else:
            return self.pcap

    def Data(self) -> tuple:
        for pkthdr, data in self._waper_data():
            dcopy_data = copy.deepcopy(data)
            # 空数据处理
            if pkthdr is None:
                continue
            # print(f"pkthdr: {pkthdr}, data: {data}")

            # 过滤掉ARP等其他报文
            ethHeader = IPParser.ParseEthHeader(dcopy_data)
            if ethHeader["typ_eth"] != "IPv4":
                # print(f"Unmatch ethernet type: IPv4, {1}")
                continue
            ipHeader = IPParser.ParseIpHeader(dcopy_data, True)
            # 过滤IP版本不是4的报文
            if ipHeader["ip_version"] != 4:
                print(f"Abandon other version")
                continue
            # 过滤GVCP报文
            if ipHeader["dst_ip"] == "255.255.255.255":
                print("Abandon GVCP")
                continue
            # 过滤TTL=1报文
            if ipHeader["TTL"] == 1:
                print("Abandon TTL=1")
                continue
            # 过滤自己端口发出的数据
            # if ipHeader["src_ip"] == self.channel:
            #     print(f"Abandon self-data")
            #     continue
            # 过滤掉除TCP/UDP之外的报文
            if ipHeader["protocol"] not in [6, 17]:
                print(f"Abandon other protocol")
                continue
            transHeader = IPParser.ParseTransHeader(dcopy_data, True)
            # # 过滤掉纯ACK报文
            # if transHeader["flags"] == b'\x00\x10':
            #     print("Abandon ACK")
            #     continue
            # 过滤掉非有效数据报文
            if (transHeader["flags"][-1] & 0x0f) >> 3 != 1:
                print(f"Abandon is not push the data. {transHeader['flags']}")
                continue
            # 过滤多播DNS(MDNS)报文
            if ipHeader["dst_ip"] == "224.0.0.251" and transHeader["dst_port"] == 5353:
                print("Abandon MDNS")
                continue
            # 过滤NBNS报文
            if transHeader["src_port"] == 137 and transHeader["dst_port"] == 137:
                print("Abandon NBNS")
                continue
            payload = dcopy_data[IPParser.ETH_MAC_LENGTH + IPParser.IP_HEADER_LENGTH + IPParser.TRANS_CTRL_LENGTH:]
            yield (pkthdr.ts, ipHeader["src_ip"], transHeader["src_port"], ipHeader["dst_ip"], transHeader["dst_port"],
                   payload)

    def Destroy(self) -> None:
        if not self.pcap:
            return
        try:
            self.pcap.close()
        except Exception:
            pass
        self.pcap = None


class DoIPMsg:
    def __init__(self, message=None, verbose=False):
        self.UpdateMsg(message, verbose)

    def UpdateMsg(self, message=None, verbose=False):
        if not message:
            self.messageString = None
            self.protcolVersion = self.inverseProtocolVersion = None
            self.payloadType = self.payloadLength = None
            self.sourceAddress = self.targetAddress = None
            self.payload = None
            self.isUDS = False
        else:
            self.messageString = message
            self.protcolVersion = message[0:2]
            self.inverseProtocolVersion = message[2:4]
            self.payloadType = message[4:8]
            self.payloadLength = message[8:16]
            self.sourceAddress = message[16:20]
            if self.payloadType == DOIP_ROUTING_ACTIVATION_REQUEST:
                self.targetAddress = None
            else:
                self.targetAddress = message[20:24]

            if self.payloadType == DOIP_DIAGNOSTIC_MESSAGE:
                self.isUDS = True
                self.payload = message[24:len(message)]
            else:
                self.payload = message[24:len(message) - len(ASRBISO)]
                self.isUDS = False
            if verbose:
                print(str(message))
                self.PrintMessage()

    def PrintMessage(self):
        print("Protocol Version 		: " + str(self.protcolVersion))
        print("Inv. Protocol Version 	: " + str(self.inverseProtocolVersion))
        print("Payload Type 			: " + str(self.payloadType))
        print("Payload Type Description : " + str(self.DecodePayloadType(self.payloadType)))
        print("Payload Length 			: " + str(self.payloadLength))
        print("Source Address 			: " + str(self.sourceAddress))
        print("Target Address 			: " + str(self.targetAddress))
        print("Payload 					: " + str(self.payload))
        print("")

    def DecodePayloadType(self, payloadType=None):
        if payloadType is None:
            payloadType = self.payloadType
        return payloadTypeDescription.get(int(payloadType), "Invalid or unregistered diagnostic payload type")


if __name__ == '__main__':
    ec = EthCapture("192.168.20.2")
    ec.Start()
